#ifndef _HYBRID_LOOKUP_ACCELERATOR_H_
#define _HYBRID_LOOKUP_ACCELERATOR_H_

//hash-kd-tree hybrid
#include <vector>
#include <memory>
#include "parallel.h"
#include "kdtree.h"

template<typename T>
class HybridLookupAccelerator{
	enum HashCellType{HashCellKDTree = 0, HashCellArray};
	struct KDTree{
		KDTree(const std::vector<size_t>& indices, const std::vector<T>& all_photons)
		{
			Build(indices, all_photons);
		}
		template<typename P, typename ProcData, typename Proc, typename DistT>
		void Lookup(const P& p, ProcData& proc_data, const Proc& proc, DistT& dist, const std::vector<T>& all_photons)const;
		~KDTree(){
			index_kd_tree_destroy(&index_kdtree_);
		}
	private:
		void Build(const std::vector<size_t>& indices, const std::vector<T>& all_photons);
		typedef kd_tree_t(size_t) index_kdtree_t;
		index_kdtree_t index_kdtree_;
	};
	struct HashCell{
		HashCell():hash_cell_type_(HashCellArray){
			index_ = NULL;
		}
		void AddPhotonIndex(size_t idx){
			assert(hash_cell_type_ == HashCellArray);

			if (index_ == NULL)
				index_ = new std::vector<size_t>();
			index_->push_back(idx);
		}

		void TryBuildKDTree(size_t min_count, const std::vector<T>& all_photons);
		
		template<typename P, typename ProcData, typename Proc, typename DistT>
		void Lookup(const P& p, ProcData& proc_data, const Proc& proc, DistT& dist, const std::vector<T>& all_photons)const;
		~HashCell(){
			if (hash_cell_type_ == HashCellArray)
				delete index_;
			else
				delete kd_tree_;
		}

	private:
		HashCellType hash_cell_type_;
		union{
			KDTree* kd_tree_;
			std::vector<size_t>* index_;
		};
	};

	std::shared_ptr<std::vector<T> > all_photons_;
	HashCell* hash_cells_;
	size_t hash_grid_size_;
	bbox_t bbox_;
	float inv_cell_size_;
private:
	unsigned int hash(const int ix, const int iy, const int iz)const {
		return (unsigned int)((ix * 73856093) ^ (iy * 19349663) ^ (iz * 83492791)) % hash_grid_size_;
	}
public:
	HybridLookupAccelerator(size_t nphoton_lookup, float max_radius2, std::shared_ptr<std::vector<T> > photons);
	HybridLookupAccelerator(std::shared_ptr<std::vector<T> > photons);//kd_tree only

	template<typename P, typename ProcData, typename Proc, typename DistT>
	void Lookup(const P& p,  ProcData& proc_data, const Proc& proc, DistT& dist)const;
	bool Empty()const{return hash_grid_size_ == 0;}
	size_t Size()const{return all_photons_->size();}
	~HybridLookupAccelerator()
	{delete []hash_cells_;}
};

template<typename T>
void HybridLookupAccelerator<T>::HashCell::TryBuildKDTree(size_t min_count, const std::vector<T>& all_photons)
{
	assert(hash_cell_type_ == HashCellArray);
	if(index_ && index_->size() > min_count)
	{
		hash_cell_type_ = HashCellKDTree;
		std::vector<size_t> *index = index_;
		kd_tree_ = new KDTree(*index,all_photons);
		delete index;
	}
}
template<typename T>
template<typename P, typename ProcData, typename Proc, typename DistT>
void HybridLookupAccelerator<T>::HashCell::Lookup(const P& p, ProcData& proc_data, const Proc& proc, DistT& max_dist2,const std::vector<T>& all_photons)const
{
	DistT debug_max_dist2 = max_dist2;
	if(hash_cell_type_ == HashCellArray && index_)
	{
		for(size_t i = 0;i < index_->size(); ++i)
		{
			float dist2 = distance_squared(all_photons[(*index_)[i]].p, p);
			if(dist2 <= max_dist2)
			{
				(proc)(&(proc_data), const_cast<T*>(&all_photons[(*index_)[i]]), dist2, &(max_dist2));
			}
		}
	}
	else if(hash_cell_type_ == HashCellKDTree && kd_tree_)
	{
		kd_tree_->Lookup(p,proc_data,proc,max_dist2,all_photons);
	}
}
	
	
//kd-tree only
template<typename T>
HybridLookupAccelerator<T>::HybridLookupAccelerator(std::shared_ptr< std::vector<T> > photons)
{
	all_photons_ = photons;
	hash_grid_size_ = 1;
	hash_cells_ = new HashCell[hash_grid_size_];	
	bbox_init(&bbox_);
	for(size_t i = 0;i < all_photons_->size(); ++i)
	{
		hash_cells_[0].AddPhotonIndex(i);
	}
	hash_cells_[0].TryBuildKDTree(0,*all_photons_);
}
template<typename T>
HybridLookupAccelerator<T>::HybridLookupAccelerator(size_t nlookup, float max_radius2, std::shared_ptr<std::vector<T> > photons)
{
	const float photon_radius = sqrtf(max_radius2);
	const float cell_size = photon_radius * 2.f;
	inv_cell_size_ = 1.f/cell_size;

	all_photons_ = photons;

	hash_grid_size_ = all_photons_->size();
	hash_cells_ = new HashCell[hash_grid_size_];	

	bbox_init(&bbox_);
	for(size_t i = 0;i < all_photons_->size(); ++i)
	{
		bbox_union_with_point(&bbox_, & (*all_photons_)[i].p);
	}
	for(size_t i = 0;i < all_photons_->size(); ++i)
	{
		vector3f_t rad;
		vinit(rad,photon_radius, photon_radius, photon_radius);
		vector3f_t b_min,b_max;
		vsub(b_min, (*all_photons_)[i].p, rad);
		vsub(b_min, b_min, bbox_.pmin);
		vsmul(b_min, inv_cell_size_, b_min);

		vadd(b_max, (*all_photons_)[i].p, rad);
		vsub(b_max, b_max, bbox_.pmin);
		vsmul(b_max, inv_cell_size_, b_max);

		for(int iz = abs(int(b_min.z)); iz <= abs(int(b_max.z)); ++iz)
		{
			for(int iy = abs(int(b_min.y)); iy <= abs(int(b_max.y));++iy)
			{
				for(int ix = abs(int(b_min.x)); ix <= abs(int(b_max.x)); ++ix)
				{
					unsigned int hv = hash(ix,iy,iz);//if collision may add twice
					hash_cells_[hv].AddPhotonIndex(i);
				}
			}
		}
	}

	/*
	for(size_t i = 0;i < hash_grid_size_; ++i)
	{
		hash_cells_[i].TryBuildKDTree(nlookup,*all_photons_);
	}

	*/
	tbb::parallel_for(tbb::blocked_range<size_t>(0,hash_grid_size_),
		[&](const tbb::blocked_range<size_t>& r){
			for(size_t i = r.begin(); i != r.end(); ++i)
				hash_cells_[i].TryBuildKDTree(nlookup,*all_photons_);
	}
	);
}

template<typename T>
template<typename P, typename ProcData, typename Proc, typename DistT>
void HybridLookupAccelerator<T>::Lookup(const P& p, ProcData& proc_data, const Proc& proc, DistT& dist)const
{
	vector3f_t hh ;
	vsub(hh, p, bbox_.pmin);
	vsmul(hh, inv_cell_size_, hh);
	const int ix = abs(int(hh.x));
	const int iy = abs(int(hh.y));
	const int iz = abs(int(hh.z));
	const HashCell& hc = hash_cells_[hash(ix,iy,iz)];
	hc.Lookup(p,proc_data, proc, dist, *all_photons_);
}
template<typename T>
void HybridLookupAccelerator<T>::KDTree::Build(const std::vector<size_t>& indices, const std::vector<T>& all_photons)
{
	index_kd_tree_init(&index_kdtree_,indices,all_photons);
}
template<typename T>
template<typename P, typename ProcData, typename Proc, typename DistT>
void HybridLookupAccelerator<T>::KDTree::Lookup(const P& p, ProcData& proc_data, const Proc& proc, DistT& dist,const std::vector<T>& all_photons)const
{
	std::vector<T>& all_photons_data = const_cast<std::vector<T>& >(all_photons);
	index_kd_tree_lookup(index_kdtree_,p,(&proc_data),proc,dist,all_photons_data);
}


#endif
